#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# QQ: 34538980@qq.com
# Jekkay Hu, 2013.5.5
# ///////////////////////////////////////////////////////////////////
#                            _ooOoo_                               //
#                           o8888888o                              //
#                           88" . "88                              //
#                           (| ^_^ |)                              //
#                           O\  =  /O                              //
#                        ____/`---'\____                           //
#                      .'  \\|     |//  `.                         //
#                     /  \\|||  :  |||//  \                        //
#                    /  _||||| -:- |||||-  \                       //
#                    |   | \\\  -  /// |   |                       //
#                    | \_|  ''\---/''  |   |                       //
#                    \  .-\__  `-`  ___/-. /                       //
#                  ___`. .'  /--.--\  `. . ___                     //
#                ."" '<  `.___\_<|>_/___.'  >'"".                  //
#              | | :  `- \`.;`\ _ /`;.`/ - ` : | |                 //
#              \  \ `-.   \_ __\ /__ _/   .-` /  /                 //
#        ========`-.____`-.___\_____/___.-`____.-'========         //
#                             `=---='                              //
#        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^        //
#             佛祖保佑       永无BUG        运行正常                   //
#     -------------------------------------------------------      //
#               QQ: 34538980@qq.com                                //
#               博客: http://www.easysb.cn                          //
#               Jekkay Hu, 2013.5.5                                //
# ///////////////////////////////////////////////////////////////////
#
# 从数据库中提取待扫描的域名，然后加入到队里中，等待扫描，扫描线程池会从中获取目标然后扫描
#
import time
import logging
from conf.const import SCAN_STATE_INIT, SCAN_STATE_QUEUE
from scan.base_db_op_thread import BaseDBOpThread

SQL_QUERY_READY_DOMAIN = 'SELECT * FROM `domain` WHERE scan = {} LIMIT %(limit)s'.format(SCAN_STATE_INIT)
SQL_UPDATE_QUEUE_DOMAIN = 'UPDATE `domain` SET scan=%s WHERE id IN ({})' % SCAN_STATE_QUEUE


class DomainSupplierThread(BaseDBOpThread):
    def __init__(self, domain_queue):
        super(DomainSupplierThread, self).__init__()
        self.domain_queue = domain_queue
        self.threshold = 10

    def work_loop(self):
        if self.domain_queue.qsize() < self.threshold:
            size = self._add_domain_to_queue(self.threshold - self.domain_queue.qsize())
            if size <= 0:
                # 出现不够就休息10秒
                time.sleep(10)
                self._close_db_connect()
        # 休息100毫秒
        time.sleep(0.1)

    def _add_domain_to_queue(self, count):
        size = 0
        while count > size:
            result = self.db.query(SQL_QUERY_READY_DOMAIN, {'limit': min(30, count - size)})
            if not result:
                logging.warning("not enough ready domains %s/%s" % (size, count))
                return size
            self._mark_domain_status_queue(list(map(lambda x: x['id'], result)))
            size += len(result)
            for item in result:
                self.domain_queue.put(item)
        return size

    def _mark_domain_status_queue(self, id_list):
        if not id_list:
            return
        sql = SQL_UPDATE_QUEUE_DOMAIN.format(','.join(map(lambda x: '"%s"' % x, id_list)))
        self.db.execute(sql)
